Before use this notebook¶
This notebook is based on 'https://colab.research.google.com/github/roboflow/notebooks/blob/main/notebooks/train-yolov10-object-detection-on-custom-dataset.ipynb'
Some modifications were made to train custom data.
import os
HOME = os.getcwd()
!nvidia-smi
!lscpu
!free
Sat Feb 8 00:26:20 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.120 Driver Version: 560.94 CUDA Version: 12.6 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA GeForce RTX 3060 Ti On | 00000000:07:00.0 On | N/A |
| 0% 32C P8 17W / 200W | 895MiB / 8192MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| 0 N/A N/A 815 C /python3.10 N/A |
+-----------------------------------------------------------------------------------------+
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 48 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 6
On-line CPU(s) list: 0-5
Vendor ID: AuthenticAMD
Model name: AMD Ryzen 5 5600 6-Core Processor
CPU family: 25
Model: 33
Thread(s) per core: 1
Core(s) per socket: 6
Socket(s): 1
Stepping: 2
BogoMIPS: 6986.95
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge m
ca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall
nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep
_good nopl tsc_reliable nonstop_tsc cpuid extd_apicid
pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcn
t aes xsave avx f16c rdrand hypervisor lahf_lm cmp_leg
acy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch
osvw topoext perfctr_core ssbd ibrs ibpb stibp vmmcal
l fsgsbase bmi1 avx2 smep bmi2 erms invpcid rdseed adx
smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 x
saves clzero xsaveerptr arat npt nrip_save tsc_scale v
mcb_clean flushbyasid decodeassists pausefilter pfthre
shold v_vmsave_vmload umip vaes vpclmulqdq rdpid fsrm
Virtualization features:
Virtualization: AMD-V
Hypervisor vendor: Microsoft
Virtualization type: full
Caches (sum of all):
L1d: 192 KiB (6 instances)
L1i: 192 KiB (6 instances)
L2: 3 MiB (6 instances)
L3: 32 MiB (1 instance)
Vulnerabilities:
Gather data sampling: Not affected
Itlb multihit: Not affected
L1tf: Not affected
Mds: Not affected
Meltdown: Not affected
Mmio stale data: Not affected
Reg file data sampling: Not affected
Retbleed: Not affected
Spec rstack overflow: Mitigation; safe RET, no microcode
Spec store bypass: Mitigation; Speculative Store Bypass disabled via prct
l and seccomp
Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointe
r sanitization
Spectre v2: Mitigation; Retpolines; IBPB conditional; IBRS_FW; STI
BP disabled; RSB filling; PBRSB-eIBRS Not affected; BH
I Not affected
Srbds: Not affected
Tsx async abort: Not affected
total used free shared buff/cache available
Mem: 8137156 1535784 6072172 82260 529200 6280860
Swap: 4194304 0 4194304
Install YOLOv10¶
NOTE: Currently, YOLOv10 does not have its own PyPI package. Therefore, we need to install the code from the source.
!pip install -q git+https://github.com/THU-MIG/yolov10.git
NOTE: We will also install two additional packages: roboflow to download the dataset from Roboflow Universe, which we will use to train our model, and supervision, which we will use for visualizing the results.
!pip install supervision roboflow
Download pre-trained weights¶
NOTE: YOLOv10 provides weight files pre-trained on the COCO dataset in various sizes. Let's download them.
!mkdir -p {HOME}/weights
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10n.pt
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10s.pt
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10m.pt
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10b.pt
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10x.pt
!wget -P {HOME}/weights -q https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10l.pt
!ls -lh {HOME}/weights
Custom Training¶
%cd {HOME}
# Before run train command, Print it
!echo yolo task=detect mode=train epochs=10 batch=32 plots=True model={HOME}/weights/yolov10n.pt data={HOME}/data/data.yaml
# Run training!
!yolo task=detect mode=train epochs=10 batch=32 plots=True model={HOME}/weights/yolov10n.pt data={HOME}/data/data.yaml
NOTE: change to last train directory
last_train_dir=!ls -d {HOME}/runs/detect/train* | sort -V | tail -n 1
if len(last_train_dir):
print(last_train_dir[0])
global TRAIN_DIR
TRAIN_DIR=last_train_dir[0]
else:
print("no train folder!")
NOTE: Let's display train result
from IPython.display import Image
Image(filename=f'{TRAIN_DIR}/confusion_matrix.png', width=600)
Image(filename=f'{TRAIN_DIR}/results.png', width=600)
Inference with Custom Model¶
NOTE: Let's start by loading our newly trained model.
from ultralytics import YOLO
import supervision as sv
dataset = sv.DetectionDataset.from_yolo(
images_directory_path=f"{HOME}/data/images/valid",
annotations_directory_path=f"{HOME}/data/labels/valid",
data_yaml_path=f"{HOME}/data/data.yaml"
)
bounding_box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
NOTE: Let's display images from our validation set and visualize the results.
model = YOLO(f'{TRAIN_DIR}/weights/best.pt')
for path, image, annotation in dataset:
results = model(source=image, conf=0.5)[0]
detections = sv.Detections.from_ultralytics(results)
annotated_image = bounding_box_annotator.annotate(scene=image, detections=detections)
annotated_image = label_annotator.annotate(scene=image, detections=detections)
sv.plot_image(annotated_image)
0: 384x640 4 _flawless_amethysts, 1 _flawless_diamond, 3 _flawless_saphires, 3 _perfect_diamonds, 3 _perfect_emeralds, 5 _perfect_rubys, 3 _perfect_saphires, 1 _perfect_skull, 26.5ms Speed: 2.5ms preprocess, 26.5ms inference, 1.8ms postprocess per image at shape (1, 3, 384, 640)
0: 384x640 3 _flawless_amethysts, 2 _flawless_emeralds, 1 _flawless_ruby, 1 _flawless_saphire, 6 _flawless_skulls, 3 _perfect_amethysts, 2 _perfect_diamonds, 4 _perfect_emeralds, 2 _perfect_rubys, 2 _perfect_saphires, 8 _perfect_skulls, 35.9ms Speed: 2.7ms preprocess, 35.9ms inference, 2.3ms postprocess per image at shape (1, 3, 384, 640)
0: 384x640 85 _perfect_amethysts, 14 _perfect_skulls, 24.4ms Speed: 2.7ms preprocess, 24.4ms inference, 2.0ms postprocess per image at shape (1, 3, 384, 640)
0: 384x640 2 _flawless_diamonds, 2 _flawless_emeralds, 3 _flawless_rubys, 1 _flawless_saphire, 1 _flawless_skull, 4 _perfect_diamonds, 23 _perfect_rubys, 1 _perfect_saphire, 25.8ms Speed: 2.6ms preprocess, 25.8ms inference, 1.9ms postprocess per image at shape (1, 3, 384, 640)